import pandas as pd
import numpy as np
import rasterio
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_squared_error
from sklearn.preprocessing import StandardScaler
from joblib import dump
import os
import warnings
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize

# 忽略警告
warnings.filterwarnings("ignore")

# 设置工作目录
data_dir = r"D:\data1"

# 读取公益林数据
forest_data = pd.read_excel(os.path.join(data_dir, "2021公益林特征数据（2024.12）.xls"))

# 检查数据框的列名
print("数据框列名:", forest_data.columns.tolist())

# 提取2015年NDVI栅格数据
ndvi_file = os.path.join(data_dir, "hunan_ndvi_2015.tif")  # 修改为2015年NDVI文件
with rasterio.open(ndvi_file) as src:
    ndvi_profile = src.profile
    ndvi_data = src.read(1)  # 假设是单波段数据

# 将经纬度转换为栅格坐标并提取NDVI值
def extract_ndvi(row, src):
    # 获取经纬度
    lon, lat = row['纵坐标'], row['横坐标']
    # 将经纬度转换为栅格坐标
    try:
        row_idx, col_idx = src.index(lon, lat)
        # 提取NDVI值
        if 0 <= row_idx < ndvi_data.shape[0] and 0 <= col_idx < ndvi_data.shape[1]:
            return ndvi_data[row_idx, col_idx]
        else:
            return np.nan
    except:
        return np.nan

# 为每个样地提取NDVI值
forest_data['extracted_ndvi'] = forest_data.apply(lambda row: extract_ndvi(row, src), axis=1)

# 数据清洗：去除缺失值
forest_data = forest_data.dropna(subset=['extracted_ndvi', '蓄积量'])

# 特征选择：选择与生物量相关的变量
# 确保列名匹配数据框中的实际列名
features = ['extracted_ndvi', '海拔（m', 'NDVI', 'EVI', 'GNDVI', 'MSAVI', 'SAVI', 'WDVI', 'RVI', 'DVI', 'NDVI45']
X = forest_data[features]
y = forest_data['蓄积量']  # 蓄积量作为目标变量

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 特征标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 使用随机森林回归模型
model = RandomForestRegressor(
    n_estimators=100,
    max_depth=10,
    min_samples_split=5,
    min_samples_leaf=2,
    random_state=42,
    n_jobs=-1
)
model.fit(X_train_scaled, y_train)

# 模型评估
y_pred = model.predict(X_test_scaled)
r2 = r2_score(y_test, y_pred)
mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)

print(f"模型R²: {r2:.4f}")
print(f"均方误差 (MSE): {mse:.4f}")
print(f"均方根误差 (RMSE): {rmse:.4f}")

# 保存模型
model_path = os.path.join(data_dir, "biomass_model_2015.joblib")  # 修改模型保存路径
dump(model, model_path)
print(f"模型已保存到: {model_path}")

# 应用模型到整个2015年NDVI栅格数据
# 创建预测数据（示例：仅使用extracted_ndvi）
ndvi_flat = ndvi_data.flatten()
valid_mask = ~np.isnan(ndvi_flat)

# 创建预测数据（示例：仅使用extracted_ndvi）
X_grid = np.zeros((len(ndvi_flat), len(features)))
X_grid[:, 0] = ndvi_flat  # 第一列是extracted_ndvi
# 其他特征需要根据实际情况填充，这里用0填充示例
X_grid[~valid_mask] = 0

# 标准化预测数据
X_grid_scaled = scaler.transform(X_grid)

# 预测生物量
biomass_pred = model.predict(X_grid_scaled)
biomass_pred = biomass_pred.reshape(ndvi_data.shape)

# 保存预测结果为GeoTIFF
output_file = os.path.join(data_dir, "hunan_biomass_2015.tif")  # 修改输出文件名
profile = ndvi_profile.copy()
profile.update(dtype=rasterio.float32, nodata=0)
with rasterio.open(output_file, 'w', **profile) as dst:
    dst.write(biomass_pred.astype(np.float32), 1)

print(f"生物量预测结果已保存到: {output_file}")

# 上色处理并保存为PNG
plt.figure(figsize=(12, 10))
cmap = plt.cm.YlOrRd  # 使用YlOrRd颜色映射，也可以选择其他映射如plt.cm.viridis
norm = Normalize(vmin=np.nanmin(biomass_pred), vmax=np.nanmax(biomass_pred))

# 获取栅格数据的地理范围
bounds = rasterio.open(output_file).bounds
extent = [bounds.left, bounds.right, bounds.bottom, bounds.top]

# 绘制上色后的图像
im = plt.imshow(biomass_pred, cmap=cmap, norm=norm, extent=extent)
plt.colorbar(im, label='Biomass Value')
plt.title('2015 Hunan Biomass Visualization')
plt.xlabel('Longitude')
plt.ylabel('Latitude')

# 保存为PNG格式
colored_png_file = os.path.join(data_dir, "hunan_biomass_2015_colored.png")  # 修改PNG文件名
plt.savefig(colored_png_file, dpi=300, bbox_inches='tight')
print(f"上色后的生物量预测结果已保存为PNG: {colored_png_file}")

# 显示图像
plt.show()
